#ifndef _EKF_HPP
#define _EKF_HPP

#include <iostream>
#include <fstream>
#include <memory>
#include <eigen3/Eigen/Dense>
#include "nav_base.hpp"
#include "mec.hpp"

#define N 16

class Data{
public:
    ImuData imudata;
    OdomData odomdata;
};
using DataPtr=std::shared_ptr<Data>;

class EKF:NavBase{
private:
    Eigen::Matrix<double, N, 1> delta_x_;   // 误差状态：δp，δv，phi，bg，ba，k(里程计比例因子)
    Eigen::Vector3d delta_z_;               // 速度误差观测值
    DataPtr data_curr_;
    DataPtr data_last_;
    StatePtr state_;
    Mec mec_;
    ImuModel imumodel_;

    Eigen::Matrix<double, N, N> Pmat_;

    Eigen::Matrix<double, N, N> PHImat_;
    Eigen::Matrix<double, N, N> Qmat_;

    Eigen::Matrix<double, N, N> Gmat_;      // 噪声输入映射矩阵 noise-input mapping matrix
    Eigen::Matrix<double, N, N> Gmat_last_;
    Eigen::Matrix<double, N, N> qmat_;
    Eigen::Matrix<double, N, N> qmat_last_; 

    Eigen::Matrix<double, N, 3> Kmat_;

    Eigen::Matrix3d             Rmat_; // 里程计观测噪声协方差矩阵
    
    std::ofstream result_fs; 
    

private:
    /**
     * @brief 预测
     * 
     * @param dt 
     */
    void KalmanPredict(double dt){

        auto I_33 = Eigen::Matrix3d::Identity();
        auto Zero_33 = Eigen::Matrix3d::Zero();

        // 计算状态转移矩阵Φ
        PHImat_ = Eigen::Matrix<double, N, N>::Identity();

        PHImat_.block<3, 3>(0, 3) = I_33 * dt;

        Eigen::Vector3d fn = state_->c_bn * data_curr_->imudata.accel_mpss;
        PHImat_.block<3, 3>(3, 6) = getMX(fn) * dt;
        PHImat_.block<3, 3>(3, 12) = state_->c_bn * dt;
        // 有错，少加了一句：
        PHImat_.block<3, 3>(6, 6) = I_33 - getMX(data_curr_->imudata.gyro_rps) * dt;
        PHImat_.block<3, 3>(6, 9) = -state_->c_bn * dt;

        // 计算状态转移噪声协方差矩阵Q
        double item[] = {   
                            imumodel_.VRW * imumodel_.VRW,
                            imumodel_.ARW * imumodel_.ARW,
                            2 * imumodel_.gyro_bias_std * imumodel_.gyro_bias_std / imumodel_.gyro_bias_corr_time,
                            2 * imumodel_.accel_bias_std * imumodel_.accel_bias_std / imumodel_.accel_bias_corr_time,
                            imumodel_.odom_scale_std * imumodel_.odom_scale_std,
                            0,
                        };
        Eigen::Matrix3d mat[6];
        qmat_ = Eigen::Matrix<double, N, N>::Zero(); // Q阵的连续时间形式Q(t)
        for(int i = 0; i < 4; i++){
            mat[i] = Eigen::Vector3d(item[i], item[i], item[i]).asDiagonal();
            qmat_.block<3, 3>(3 * i,  3 * i) = mat[i];
        }
        qmat_(12, 12) = item[4];

        auto Cbn = state_->c_bn;
        Gmat_ = Eigen::Matrix<double, N, N>::Zero();
        Gmat_.block<3, 3>( 3, 0) = Cbn;
        Gmat_.block<3, 3>( 6, 3) = Cbn;
        Gmat_.block<3, 3>( 9, 6) = I_33;
        Gmat_.block<3, 3>(12, 9) = I_33;
        Gmat_(15, 12) = 1;
        
        // 梯形积分
        Qmat_ = 0.5 * (PHImat_ * Gmat_last_ * qmat_ * Gmat_last_.transpose() * PHImat_.transpose() 
                + Gmat_ * qmat_ * Gmat_.transpose()) 
                * dt;

        Gmat_last_ = Gmat_;

        // 状态更新
        delta_x_ = PHImat_ * delta_x_;
        Pmat_ = PHImat_ * Pmat_ * PHImat_.transpose() + Qmat_;
    }
    /**
     * @brief 校正
     * 
     * @param dt 
     */
    void KalmanCorrect(double dt){
        auto Cbn     = state_->c_bn;
        auto I_33    = Eigen::Matrix3d::Identity();
        auto Zero_33 = Eigen::Matrix3d::Zero();

        Eigen::Matrix3d Cnb = Cbn.transpose();

        delta_z_ = state_->c_bv * Cnb * state_->vel 
                - data_curr_->odomdata.vel;

        // 观测矩阵
        Eigen::Matrix<double, 3, N> Hmat;
        Hmat << Zero_33, 
                state_->c_bv * Cnb, 
                -state_->c_bv * Cnb * getMX(state_->vel), 
                -state_->c_bv * getMX(state_->odom_lever_arm), 
                Zero_33,
                Eigen::Matrix<double, 3, 1>(-data_curr_->odomdata.vel[0],
                                            -data_curr_->odomdata.vel[0],
                                            -data_curr_->odomdata.vel[0]);

        // 卡尔曼增益
        Kmat_ = Pmat_ * Hmat.transpose() * (Hmat * Pmat_ * Hmat.transpose() + Rmat_).inverse();

        // 后验状态
        delta_x_ = delta_x_ + Kmat_ * (delta_z_ - Hmat * delta_x_);
        // 后验误差协方差矩阵
        auto I_NN = Eigen::Matrix<double, N, N>::Identity();
        Pmat_ = (I_NN - Kmat_ * Hmat)
                * Pmat_
                * (I_NN - Kmat_ * Hmat).transpose()
                + Kmat_ * Rmat_ * Kmat_.transpose();
    }

    /**
     * @brief 传感器原始数据修正
     * 
     */
    void DataErrCompensate(){
        // IMU误差补偿
        for(int i = 0; i < 3; i++){
            data_curr_->imudata.gyro_rps[i]  -= state_->gyro_bias[i];
            data_curr_->imudata.accel_mpss[i] -= state_->accel_bias[i];
        }

        // 里程计误差补偿
        data_curr_->odomdata.vel[0]  /= (1 + state_->odom_scale_factor);
    }

    /**
     * @brief 状态误差修正
     * 
     * @param dt 
     */
    void StateCorrect(double dt){
        // 位置修正
        state_->pos -= delta_x_.block<3, 1>(0, 0);

        // 速度修正
        state_->vel -= delta_x_.block<3, 1>(3, 0);

        // 姿态修正
        Eigen::Vector3d rotvec = delta_x_.block<3, 1>(6, 0);
        Eigen::Vector4d temp;        
        temp << sin(0.5 * rotvec.norm()) / (0.5 * rotvec.norm()) * 0.5 * rotvec, 
                cos(0.5 * rotvec.norm());
        Eigen::Quaterniond q_np(temp);
        state_->q_bn = q_np * state_->q_bn;
        state_->q_bn.normalize();  

        // 零偏、比例因子修正
        for(int i = 0; i < 3; i++){
            state_->gyro_bias[i]    += delta_x_[ 9 + i]; 
            state_->accel_bias[i]   += delta_x_[12 + i];
        }
        state_->odom_scale_factor += delta_x_[15];


        // 保存结果
        result_fs << state_->timestamp 
            << " " << state_->gyro_bias[0] << " " << state_->gyro_bias[1] << " " << state_->gyro_bias[2]
            << " " << state_->accel_bias[0] << " " << state_->accel_bias[1] << " " << state_->accel_bias[2]
            << " " << state_->odom_scale_factor
            << " " << delta_z_[0] << " " << delta_z_[1] << " " << delta_z_[2]
            << " " << delta_x_[0] << " " << delta_x_[1] << " " << delta_x_[2]
            << " " << std::endl;

        
        // 状态清零
        delta_x_ << 0, 0, 0,
                    0, 0, 0, 
                    0, 0, 0, 
                    0, 0, 0, 
                    0, 0, 0, 0;
    }

public:
    EKF()=default;
    ~EKF()=default;

    void initialize(StatePtr& state, DataPtr& data_raw, ImuModel& imumodel){
        state_ = state;
        mec_ = Mec(state);
        imumodel_ = imumodel;

        // 将初始欧拉角转换为四元数
        euler2quat(state_->att, state_->q_bn);

        // 将安装角转换为旋转矩阵
        // 注：不能用 a = a.transpose()，因为 a 的内存没有复制
        euler2dcm(state_->imu_mount_angle, state_->c_bv);        

        data_last_ = std::make_shared<Data>(*data_raw);

        delta_x_ << 0, 0, 0,
                    0, 0, 0, 
                    0, 0, 0, 
                    0, 0, 0, 
                    0, 0, 0, 0;
        // 初始化P阵
        Eigen::Matrix3d gb = Eigen::Vector3d(imumodel_.gyro_bias_std * imumodel_.gyro_bias_std, 
                                imumodel_.gyro_bias_std * imumodel_.gyro_bias_std, 
                                imumodel_.gyro_bias_std * imumodel_.gyro_bias_std).asDiagonal();
        Eigen::Matrix3d ab = Eigen::Vector3d(imumodel_.accel_bias_std * imumodel_.accel_bias_std, 
                                imumodel_.accel_bias_std * imumodel_.accel_bias_std, 
                                imumodel_.accel_bias_std * imumodel_.accel_bias_std).asDiagonal();

        Gmat_last_ = Eigen::Matrix<double, N, N>::Identity();

        Pmat_ = Eigen::Matrix<double, N, N>::Zero();
        Pmat_.block<3, 3>(9, 9) = gb;
        Pmat_.block<3, 3>(12, 12) = ab;
        Pmat_(15, 15) = 0; // imumodel.odom_scale_std * imumodel.odom_scale_std;

        Rmat_ = Eigen::Vector3d(
            imumodel_.odom_std * imumodel_.odom_std,
            imumodel_.odom_std * imumodel_.odom_std,
            imumodel_.odom_std * imumodel_.odom_std
        ).asDiagonal();

        result_fs = std::ofstream("/home/ghowoght/workspace/meg-odom/python/config/17floor/data-17floor-1-10/result/result_ekf_mid.txt");

    }

    void Filter(const DataPtr& data_raw){
        data_curr_ = data_raw;
        auto dt = data_curr_->imudata.timestamp - data_last_->imudata.timestamp;
        
        if(data_curr_->odomdata.valid
            && data_curr_->imudata.timestamp >= data_curr_->odomdata.timestamp
            && data_last_->imudata.timestamp <= data_curr_->odomdata.timestamp)
        {
            auto dt1 = fabs(data_curr_->imudata.timestamp - data_curr_->odomdata.timestamp);
            auto dt2 = fabs(data_last_->imudata.timestamp - data_curr_->odomdata.timestamp);

            if(dt1 > dt2){ // 距离上一时刻较近，认为里程计时间为上一时刻，则先校正再预测
            
                DataErrCompensate();
                mec_.update(data_curr_->imudata);
                KalmanPredict(dt);
                KalmanCorrect(dt);
                StateCorrect(dt);
            }
            else
            { // 否则先预测再校正
                DataErrCompensate();
                mec_.update(data_curr_->imudata);
                KalmanPredict(dt);

                KalmanCorrect(dt);
                StateCorrect(dt);
            }

            data_curr_->odomdata.valid = false;

        }
        else{
            DataErrCompensate();
            mec_.update(data_curr_->imudata);
            KalmanPredict(dt);
        }

        data_last_ = std::make_shared<Data>(*data_curr_);
    }



};

#endif
